import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from tqdm import tqdm
from get_data import load_data

# PointNet 网络定义（简化版）
class PointNet(nn.Module):
    def __init__(self, num_classes=10):
        super(PointNet, self).__init__()
        
        # MLP 层，用于特征提取
        self.mlp1 = nn.Sequential(
            nn.Conv1d(3, 64, 1),
            nn.ReLU(),
            nn.Conv1d(64, 128, 1),
            nn.ReLU(),
            nn.Conv1d(128, 1024, 1),
            nn.ReLU()
        )
        
        # 分类层
        self.fc = nn.Sequential(
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, num_classes)
        )
        
    def forward(self, x):
        x = self.mlp1(x)  # [batch_size, 1024, num_points]
        x = F.max_pool1d(x, x.size(2))  # 全局最大池化
        x = x.squeeze(2)  # [batch_size, 1024]
        x = self.fc(x)  # 分类
        return x

# 训练模型
def train_model(model, train_loader, num_epochs=20):
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    model.train()
    
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        
        for points, labels in tqdm(train_loader):
            points = points
            labels = labels
            
            optimizer.zero_grad()
            outputs = model(points)  # [batch_size, num_classes]
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            
            # 计算准确率
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}, Accuracy: {100*correct/total:.2f}%")

# 主函数
def main():
    # 设置数据目录，替换为你的本地路径
    root_dir = "C:\\Users\\maohai_pang\\Desktop\\研二(上)\\项目\\github\\datasets\\ModelNet10"  # 填写你数据集的路径
    
    # 加载训练数据
    train_loader = load_data(root_dir, split='train', batch_size=32)
    
    # 初始化模型
    model = PointNet(num_classes=10)
    
    # 训练模型
    train_model(model, train_loader)

if __name__ == "__main__":
    main()
